#include "wspack.h"
#include "string_helper.h"
#include "sha1.h"
#include "base64.h"

#define SP " "
#define EOL "\r\n"
#define MIN_FRAM_LEN 2

extern CMemPool*	_memPool;

bool read_bytes_x(const char* data, unsigned int len, char *cb, int size, int &oft)
{
	if (oft + size <= len)
	{
		memcpy(cb, (data + oft), size);
		oft += size;
		return true;
	}
	else
	{
		return false;
	}
}

CWspack::CWspack():m_r(_memPool)
{
    m_fin = 0;
    m_rsv1 = 0;
    m_rsv2 = 0;
    m_rsv3 = 0;
    m_opcode = 0;
    m_mask = 0;
    //m_length_type = 0;
    ::memset(m_masking_key, 0, sizeof(m_masking_key));
    m_payload_length = 0;
}

CWspack::~CWspack()
{
    m_r.skip(m_r.data_len());
}

int CWspack::recv_handshake(const char* data, unsigned int len)
{
    if(len > WS_MAX_HANDSHAKE_FRAME_SIZE)
        return -1;

    std::string str(data, len);
    int frame_size = fetch_hs_element(str);
    if(frame_size <= 0)
    {
        clear();
        return frame_size;
    }

    if (get_param("Upgrade") != "websocket" || get_param("Connection") != "Upgrade" ||
		get_param("Sec-WebSocket-Version") != "13" || get_param("Sec-WebSocket-Key") == "")
	{
        clear();
		return -1;
	}

    return frame_size;
}

int CWspack::pack_handshake_rsp(void)
{
	std::string magic_key = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
	std::string raw_key = get_param("Sec-WebSocket-Key") + magic_key;

	std::string sha1_key = SHA1::SHA1HashString(raw_key);
	char accept_key[128] = {0};
	Base64encode(accept_key, sha1_key.c_str(), sha1_key.length());

	std::ostringstream sstream;
	sstream << "HTTP/1.1 101 Switching Protocols" << EOL;
	sstream << "Connection: upgrade" << EOL;
	sstream << "Upgrade: websocket" << EOL;
	sstream << "Sec-WebSocket-Accept: " << accept_key << EOL;
	if (get_param("Sec-WebSocket-Protocol") != "")
	{
		sstream << "Sec-WebSocket-Protocol: chat" << EOL;
	}
	sstream << EOL;
	m_handshakeRsp = sstream.str();

	return 0;
}

int CWspack::recv_dataframe(const char* data, unsigned int len)
{
    int header_size = fetch_frame_info(data, len);
    if (header_size < 0)
    {
        clear();
        return -1;
    }
    else if (0 == header_size)
    {
        return 0;
    }
    
    if (m_payload_length + header_size > len)
	{
		return 0;
	}
    fetch_payload(data, len, header_size);

    return (m_payload_length + header_size);
}

int CWspack::fetch_frame_info(const char* data, unsigned int len)
{
    if(len < MIN_FRAM_LEN)
        return 0;

    int atf = 0;
	// FIN, opcode
	uint8_t onebyte = 0;
	if(!read_bytes_x(data, len, (char *)&onebyte, 1, atf))
        return 0;
	m_fin = onebyte >> 7;
	m_opcode = onebyte & 0x7F;

	// payload length
    if (!read_bytes_x(data, len, (char *)&onebyte, 1, atf))
        return 0;

	m_mask = onebyte >> 7 & 0x01;
	uint8_t length_type = onebyte & 0x7F;

	if (length_type < 126)
	{
		m_payload_length = length_type;
	}
	else if (length_type == 126)
	{
		/*
		uint16_t len = 0;
		input.read_bytes_x((char *)&len, 2);
		len = (len << 8) | (len >>8 & 0xFF);
		payload_length_ = len;
		*/
		uint16_t length = 0;
		uint8_t array[2] = {0};
        if(!read_bytes_x(data, len, (char *)array, 2, atf))
            return 0;
		length = uint16_t(array[0] << 8) | uint16_t(array[1]);
		m_payload_length = length;
	}
	else if (length_type == 127)
	{
		// if you don't have ntohll
		uint64_t length = 0;
		uint8_t array[8] = {0};
        if(!read_bytes_x(data, len, (char *)array, 8, atf))
            return 0;
		length = (array[0] << 56) | array[1] << 48 | array[2] << 40 | array[3] << 32 
		| array[4] << 24 | array[5] << 16 | array[6] << 8 | array[7];

		if (length > 0xFFFFFFFF)
		{
			return -1;
		}
		
	}


	// masking key
	if (m_mask == 1)
	{
        if(!read_bytes_x(data, len, (char *)m_masking_key, 4, atf))
            return 0;
	}

	// return header size
	return atf;
}

int CWspack::fetch_payload(const char* data, unsigned int len, unsigned int headlen)
{
	char real = 0;
    int atf = 0;
	if (m_mask == 1)
	{
		for (uint64_t i = 0; i < m_payload_length; i++)
		{
			read_bytes_x(data + headlen, len - headlen, &real, 1, atf);
			real = real ^ m_masking_key[i % 4];
			m_r.append(&real, 1);
		}
	}
	else
	{
		m_r.append(data + headlen, m_payload_length);
	}

	return m_payload_length;
}

int CWspack::pack_dataframe(const char* data, unsigned int len)
{
    m_fin = 1;
	uint8_t onebyte = 0;
	onebyte |= (m_fin << 7);
	onebyte |= (m_rsv1 << 6);
	onebyte |= (m_rsv2 << 5);
	onebyte |= (m_rsv3 << 4);
	onebyte |= (m_opcode & 0x0F);
	m_r.append((char *)&onebyte, 1);

	onebyte = 0;
	//set mask flag
	onebyte = onebyte | (m_mask << 7);

    m_payload_length = len;
	uint8_t length_type = get_length_type();

	if (length_type < 126)
	{
		onebyte |= m_payload_length;
		m_r.append((char *)&onebyte, 1);
	}
	else if (length_type == 126)
	{
		onebyte |= length_type;
		m_r.append((char *)&onebyte, 1);

		// also can use htons
		onebyte = (m_payload_length >> 8) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = m_payload_length & 0xFF;
		m_r.append((char *)&onebyte, 1);
	}
	else if (length_type == 127)
	{
		onebyte |= length_type;
		m_r.append((char *)&onebyte, 1);

		// also can use htonll if you have it
		onebyte = (m_payload_length >> 56) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 48) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 40) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 32) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 24) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 16) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = (m_payload_length >> 8) & 0xFF;
		m_r.append((char *)&onebyte, 1);
		onebyte = m_payload_length & 0XFF;
		m_r.append((char *)&onebyte, 1);
	}
	else
	{
		return -1;
	}

	if (m_mask == 1)
	{
        int atf = 0;
		char value = 0;
		// save masking key
		m_r.append((char *)m_masking_key, 4);
		for (uint64_t i = 0; i < m_payload_length; i++)
		{
			read_bytes_x(data, len, &value, 1, atf);
			value = value ^ m_masking_key[i % 4];
			m_r.append(&value, 1);
		}
	}
	else
	{
		m_r.append(data, m_payload_length);
	}

	return 0;
}

void CWspack::clear(void)
{
    m_params.clear();
    m_handshakeRsp.clear();
    m_r.clear();
}

int CWspack::fetch_hs_element(const std::string &msg)
{
	std::string::size_type endpos = msg.find(std::string(EOL) + EOL);
	if (endpos == std::string::npos)
	{
		return 0;
	}

	std::vector<std::string> lines;
	if (strHelper::splitStr<std::vector<std::string> >
		(lines, msg.substr(0, endpos), EOL) < 2)
	{
		return -1;
	}

	std::vector<std::string>::iterator it = lines.begin();
	while ((it != lines.end()) && strHelper::trim(*it).empty())
	{
		++it;
	};

	std::vector<std::string> reqLineParams;
	if (strHelper::splitStr<std::vector<std::string> >
		(reqLineParams, *it, SP) < 3)
	{
		return -1;
	}

	std::string mothod_ = strHelper::trim(reqLineParams[0]);
	std::string uri_ = strHelper::trim(reqLineParams[1]);
	std::string version_ = strHelper::trim(reqLineParams[2]);

	for (++it; it != lines.end(); ++it)
	{
		std::string::size_type pos = it->find_first_of(":");
		if (pos == std::string::npos)
			continue;

		std::string k = it->substr(0, pos);

		std::string v = it->substr(pos + 1);
		if (strHelper::trim(k).empty())
		{
			continue;
		}
		if (strHelper::trim(v).empty())
		{
			continue;
		}

		m_params[k] = v;
	}

	return endpos + 4;
}
const uint8_t CWspack::get_length_type()
{
	if (m_payload_length < 126)
	{
		return (uint8_t)m_payload_length;
	}
	else if (m_payload_length >= 126 && m_payload_length <= 0xFFFF)
	{
		return 126;
	}
	else
	{
		return 127;
	}
}